from typing import Dict, List, Optional

import requests


class Query:
    """A class to interact with the AtomGit paper search API.

    This class provides methods to search and retrieve papers using different criteria
    such as content, paper ID, and title.
    """

    def __init__(self, *args, **kwargs):
        """Initialize the Query instance.

        Args:
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.
        """
        self.base_url = "http://180.184.65.98:38880/atomgit"

    def query_by_content(
        self, query: str, top_k: int = 5, similarity_threshold: Optional[float] = None
    ) -> List[Dict]:
        """Search papers based on content query.

        Args:
            query (str): The search query string.
            top_k (int, optional): Maximum number of results to return. Defaults to 5.
            similarity_threshold (float, optional): Minimum similarity score to return. Defaults to None.
        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/search_papers"
        params = {"query": query, "top_k": top_k}
        try:
            response = requests.get(query_url, params=params)
            response.raise_for_status()
            results = response.json()
            if similarity_threshold:
                results = [
                    result
                    for result in results
                    if result["distance"] <= similarity_threshold
                ]
            return results
        except requests.exceptions.RequestException as e:
            print(f"Error making request: {e}")
            return []

    def query_by_id(self, paper_id: str, top_k: int = 5) -> List[Dict]:
        """Search papers based on paper ID.

        Args:
            paper_id (str): The ID of the paper to search for.
            top_k (int, optional): Maximum number of results to return. Defaults to 5.

        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/query_by_paper_id"
        params = {"paper_id": paper_id, "top_k": top_k}
        try:
            response = requests.get(query_url, params=params)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error making request: {e}")
            return []

    def query_by_title(self, title: str, top_k: int = 5) -> List[Dict]:
        """Search papers based on paper title.

        Args:
            title (str): The title of the paper to search for.
            top_k (int, optional): Maximum number of results to return. Defaults to 5.

        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/query_by_title"
        params = {"title": title, "top_k": top_k}
        try:
            response = requests.get(query_url, params=params)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error making request: {e}")
            return []


if __name__ == "__main__":
    query = Query()
    print(query.query_by_content("deep learning"))
    print(query.query_by_id("651b7dbc3fda6d7f06304579"))
    print(
        query.query_by_title(
            "SeA: Semantic Adversarial Augmentation for Last Layer Features from Unsupervised Representation Learning"
        )
    )
